import music21.chord
from music21 import *
import os
import pickle
from Note import *


def load_data(file_path):
    # load midi file using music21 library
    piece = converter.parse(file_path)

    # transpose all streams to C major. this process is to reduce the number of states
    # store the key of music before transposition.
    key = piece.analyze('key')

    if key.mode == "minor":
        i = interval.Interval(key.parallel.tonic, pitch.Pitch('C'))
    else:
        i = interval.Interval(key.tonic, pitch.Pitch('C'))

    # transpose the music using stored interval
    piece = piece.transpose(i)

    return piece


class preprocessing(object):
    def __init__(self):
        ## to collect and count unique chords, notes, octaves
        # lists that store unique chords and octaves
        self.key = None
        self.chords = []
        self.chord_octaves = []
        # lists for counting the number of times the chords and octaves appear
        self.chord_ref = []
        self.octave_ref = []
        self.chords_cnt = [0] * len(self.chord_ref)
        self.chord_octaves_cnt = [0] * len(self.octave_ref)
        # the same thing about notes
        self.notes = []
        self.note_octaves = []
        self.note_ref = []
        self.note_octave_ref = []
        self.notes_cnt = [0] * len(self.note_ref)
        self.note_octaves_cnt = [0] * len(self.note_octave_ref)

    def parsing(self, data_path):
        # load midi file
        piece = load_data(data_path)
        self.key = piece.analyze('key')
        # all_parts is list of melody and chords, each is sequence of [start time, duration, octave, pitch, velocity]
        all_parts = []
        # for all parts in midi file (most of our data have two parts, melody and chords)
        part_tuples = []
        for part in piece.iter.activeElementList:
            # check that the part is a piano song.
            # save the instrument name.
            try:
                track_name = part[0].bestName()
            except AttributeError:
                track_name = 'None'
            part_tuples.append(track_name)

            # part_tuples is sequence of [start time, duration, octave, pitch]

            for event in part._elements:
                # if chords or notes exist recursive (sometimes this happened in a complicated piano song file)
                if event.isStream:
                    _part_tuples = []
                    for i in event._elements:
                        _part_tuples = self.streaming(event, i, _part_tuples)
                    all_parts.append(_part_tuples)
                # normal case
                else:
                    part_tuples = self.streaming(part, event, part_tuples)
            if part_tuples != []:
                all_parts.append(part_tuples)
        return all_parts

    def streaming(self, part, event, part_tuples):
        # find the set of chords and octaves
        # save start time
        for y in event.contextSites():
            if y[0] is part:
                offset = y[1]
        # if current event is chord
        if getattr(event, 'isChord', None) and event.isChord:
            # chord pitch ordering
            octaves = []
            for pitch in event.pitches:
                octaves.append(pitch.octave)
            sort_idx = [i[0] for i in sorted(enumerate(event.pitchClasses), key=lambda x: x[1])]
            octaves = [x for (y, x) in sorted(zip(sort_idx, octaves))]
            # if current chord or octave is unique until now, add it to the list
            if event.orderedPitchClasses not in self.chords:
                self.chords.append(event.orderedPitchClasses)
            if octaves not in self.chord_octaves:
                self.chord_octaves.append(octaves)

        # if current event is note
        if getattr(event, 'isNote', None) and event.isNote:
            # find set of octaves and pitches of note
            if event.pitch.octave not in self.note_octaves:
                self.note_octaves.append(event.pitch.octave)
            if event.pitch.pitchClass not in self.notes:
                self.notes.append(event.pitch.pitchClass)

        # if current event is rest
        if getattr(event, 'isRest', None) and event.isRest:
            part_tuples.append([offset, event.quarterLength, 0, 0, 0])
        return part_tuples


if __name__ == "__main__":

    # print set of octaves and chords
    # a = preprocessing()
    # data_dir = './Nottingham/all/'
    # dataset = []
    # for files in os.listdir(data_dir):
    #     print(files)
    #     seq = a.parsing(data_dir + files)
    #     print('notes: ', a.notes)
    #     print('len(notes): ', len(a.notes))
    #     print('note_octaves: ', a.note_octaves)
    #     print('len(note_octaves): ', len(a.note_octaves))
    #     print('chords: ', a.chords)
    #     # with open('C:/Users/Guan Zi Xuan/seqgan-music-master/chords/chords', 'wb') as filepath:
    #     #     pickle.dump(a.chords, filepath)
    #     # print('len(chords): ', len(a.chords))
    #     # print('chord_octaves: ', a.chord_octaves)
    #     # with open('C:/Users/Guan Zi Xuan/seqgan-music-master/chords/octaves', 'wb') as filepath:
    #     #     pickle.dump(a.chord_octaves, filepath)
    #     # print('len(chord_octaves): ', len(a.chord_octaves))
    #     # print('\n')
    preprocessing_machine = preprocessing()
    path = "/Users/onlyblue/PycharmProjects/SongSimilarity/data/midi/Jazz/5 To 10 - Allan Holdsworth.mid"
    seq = preprocessing_machine.parsing(path)
    # print(preprocessing_machine.key)
    chords = []
    try:
        for item in preprocessing_machine.chords:
            print(type(item))
            note = [Note[i] for i in item]
            print(note)
            c = music21.chord.Chord(note)
            if c.commonName not in chords:
                chords.append(c.commonName)
    except:
        pass
    print(len(chords))
